package main

import (
	"flag"
	"io"
	"log"
	"net"
	"sync"
)

func main() {
	// 命令行参数
	var (
		listenAddr = flag.String("l", "127.0.0.1:8080", "本地监听地址")
		targetAddr = flag.String("t", "127.0.0.1:80", "目标转发地址")
	)
	flag.Parse()

	// 开始监听
	ln, err := net.Listen("tcp", *listenAddr)
	if err != nil {
		log.Fatalf("监听失败: %v", err)
	}
	log.Printf("开始转发: %s -> %s", *listenAddr, *targetAddr)

	for {
		localConn, err := ln.Accept()
		if err != nil {
			log.Printf("Accept 错误: %v", err)
			continue
		}
		go handleConn(localConn, *targetAddr)
	}
}

func handleConn(localConn net.Conn, target string) {
	defer localConn.Close()

	// 连接目标
	remoteConn, err := net.Dial("tcp", target)
	if err != nil {
		log.Printf("连接目标 %s 失败: %v", target, err)
		return
	}
	defer remoteConn.Close()

	// 双向复制数据
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		io.Copy(remoteConn, localConn)
		remoteConn.(*net.TCPConn).CloseWrite() // 半关闭，通知对端 EOF
		wg.Done()
	}()
	go func() {
		io.Copy(localConn, remoteConn)
		localConn.(*net.TCPConn).CloseWrite()
		wg.Done()
	}()
	wg.Wait()
}
